
from tn.processor import Processor
from tn.chinese.rules.cardinal import Cardinal
from tn.chinese.rules.char import Char
from tn.chinese.rules.date import Date
from tn.chinese.rules.fraction import Fraction
from tn.chinese.rules.math import Math
from tn.chinese.rules.measure import Measure
from tn.chinese.rules.money import Money
from tn.chinese.rules.postprocessor import PostProcessor
from tn.chinese.rules.preprocessor import PreProcessor
from tn.chinese.rules.sport import Sport
from tn.chinese.rules.time import Time
from tn.chinese.rules.whitelist import Whitelist

from pynini import cross
from pynini.lib.pynutil import add_weight, delete
from importlib_resources import files
from pynini import string_map, union
from pynini.lib.pynutil import insert

from pynini import accep

class CustomRule(Processor):
    def __init__(self, remove_ah=False):
        super().__init__(name='custom_rule')
        self.remove_ah = remove_ah
        self.build_tagger()
        self.build_verbalizer()

    def build_tagger(self):
        ah = add_weight(insert('ah: "') + accep('啊'), 1)
        ah2 = add_weight(cross("啊", "阿"), 1)
        tagger = (ah | (insert('value: "') + ah2)) + insert('"')
        self.tagger = self.add_tokens(tagger)

    def build_verbalizer(self):
        super().build_verbalizer()
        if self.remove_ah:
            verbalizer = self.delete_tokens(delete('ah: "啊"'))
        else:
            verbalizer = self.delete_tokens(
                delete('ah: \"') + accep('啊') + delete('\"'))
        self.verbalizer |= verbalizer



class Normalizer(Processor):
    def __init__(self,
                 cache_dir=None,
                 overwrite_cache=False,
                 remove_interjections=True,
                 remove_erhua=True,
                 traditional_to_simple=True,
                 remove_puncts=False,
                 full_to_half=True,
                 tag_oov=False):
        super().__init__(name='zh_normalizer')
        self.remove_interjections = remove_interjections
        self.remove_erhua = remove_erhua
        self.traditional_to_simple = traditional_to_simple
        self.remove_puncts = remove_puncts
        self.full_to_half = full_to_half
        self.tag_oov = tag_oov
        if cache_dir is None:
            cache_dir = files("tn")
        self.build_fst('zh_tn', cache_dir, overwrite_cache)

    def build_tagger(self):
        processor = PreProcessor(
            traditional_to_simple=self.traditional_to_simple).processor

        # 其他已有的规则
        date = add_weight(Date().tagger, 1.02)
        whitelist = add_weight(Whitelist().tagger, 1.03)
        sport = add_weight(Sport().tagger, 1.04)
        fraction = add_weight(Fraction().tagger, 1.05)
        measure = add_weight(Measure().tagger, 1.05)
        money = add_weight(Money().tagger, 1.05)
        time = add_weight(Time().tagger, 1.05)
        cardinal = add_weight(Cardinal().tagger, 1.06)
        math = add_weight(Math().tagger, 90)
        char = add_weight(Char().tagger, 100)
        
        # 自定义规则
        custom_rule = add_weight(CustomRule().tagger, 1.05)

        tagger = (date | whitelist | sport | fraction | measure | money | time
                  | cardinal | math | char | custom_rule).optimize()
        tagger = (processor @ tagger).star
        # 删除最后的空格
        self.tagger = tagger @ self.build_rule(delete(' '), r='[EOS]')

    def build_verbalizer(self):
        cardinal = Cardinal().verbalizer
        char = Char().verbalizer
        date = Date().verbalizer
        fraction = Fraction().verbalizer
        math = Math().verbalizer
        measure = Measure().verbalizer
        money = Money().verbalizer
        sport = Sport().verbalizer
        time = Time().verbalizer
        whitelist = Whitelist(remove_erhua=self.remove_erhua).verbalizer
        # 自定义规则
        custom_rule = CustomRule().verbalizer

        verbalizer = (cardinal | char | date | fraction | math | measure
                      | money | sport | time | whitelist | custom_rule).optimize()

        processor = PostProcessor(
            remove_interjections=self.remove_interjections,
            remove_puncts=self.remove_puncts,
            full_to_half=self.full_to_half,
            tag_oov=self.tag_oov).processor
        self.verbalizer = (verbalizer @ processor).star